from copy import copy
import re
rpm = re.compile('[+-]')
from rmpoly import Poly

class BaseFreeAlgebra:
  """abstract base class of free algebra on ring K
  subclasses:
  FreeAlgebra free algebra on a commutative ring
              with no elements with more than one term
  MFreeAlgebra free algebra on a commutative ring
              with elements with more than one term

  The objects fa of one of these classes constracy
  elements of FreeAlgebraElement

  fa.field generates the elements of the ring K
  K(0)  zero element in K
  K(1)  unit element in K

  fa.gens() noncommuting generators

  In the string representation ^ instead of ** is used 
  as exponent symbol.

  >>> from rmpoly import *
  >>> from fractions import Fraction as mpq
  >>> from free_algebra import FreeAlgebra
  >>> fa = FreeAlgebra('A,B',mpq)
  >>> A,B = fa.gens()
  >>> (A+B/2)**2
   +1/4*B^2 +1/2*B*A +1/2*A*B +A^2
  >>> rp,x = ncrgens('x',6,fa)
  >>> (A + B*x)**2
   +( +B^2)*x^2 +( +B*A +A*B)*x +( +A^2)
  """
  def __init__(self, pol_gens, field, **kwds):
    if isinstance(pol_gens,basestring):
      pol_gens = pol_gens.split(',')
    self.pol_gens = tuple(pol_gens)
    self.ngens = len(pol_gens)
    self.field = field
    self.gens_dict = {}
    self.dpol_gens = {}
    for i in range(self.ngens):
      self.gens_dict[self.pol_gens[i]] = chr(i+1)
      self.dpol_gens[chr(i+1)] = self.pol_gens[i]
    self.rules = ()
    self._pre_rules = []
    if 'parens' in kwds:
      self.parens = kwds['parens']
    else:
      self.parens = False
    str_field = str(self.field)
    if 'mpc' in str_field or 'Complex' in str_field:
      self.parens = True


  def gens(self):
    """return the list of generators"""
    a = []
    one = self.field(1)
    n = self.ngens
    for i in range(self.ngens):
      p = FreeAlgebraElement(self)
      expv = chr(i+1)
      p[expv] = one
      a.append(p)
    return a
    
  def mon_eval(self,s):
    """compute the tuple (expv,c) for a string 'c*A_0^e_0*...'
    """
    gens_dict = self.gens_dict
    field = self.field
    s = s.strip()
    if s[0].isdigit() and '*' not in s:
      return ('',field(s))
    if s[0].isdigit():
      a = s.split('*',1)
      coeff = field(a[0])
      a = a[1].split('*')
    else:
      a = s.split('*')
      coeff = field(1)
    expv = ''
  # e.g. a = ['A0^6', 'A0*A1^2*A0', 'A1^2']
  # split each element in ind,pw
    for x in a:
      t = x.split('^')
      ind = gens_dict[t[0]]
      if len(t) == 2:
        pw = int(t[1])
      else:
        pw = 1
      expv += ind*pw
    return (expv, coeff)

  def read_monom(self,s):
    """compute expv for the string 'A0^e0*...'
    """
    gens_dict = self.gens_dict
    field = self.field
    s = s.strip()
    a = s.split('*')
    expv = ''
    for x in a:
      t = x.split('^')
      ind = gens_dict[t[0]]
      if len(t) == 2:
        pw = int(t[1])
      else:
        pw = 1
      expv += ind*pw
    return expv

  def read(self, s):
    s = s.strip()
    p = FreeAlgebraElement(self)
    fa1 = self.field
    if '(' not in s:
      p = FreeAlgebraElement(p.fa)
      cp = fa1(s)
      if cp:
        p[''] = cp
      return p
    assert s[0] == '+' and s[1] == '('
    # nparens = number of '(' parenthesis - number of ')' parenthesis

    nparens = 1
    # prev_pos starting position of the subexpression; the subexpression ends
    # when nparens becomes zero.
    prev_pos=2
    ns = len(s)
    pos = 2
    while 1:
      if s[pos] == '(':
        nparens += 1
      elif s[pos] == ')':
        nparens -= 1
      if nparens == 0:
        # found subexpression
        s1 = s[prev_pos:pos]
        # evaluate the subexpression
        cp = fa1(s1)
        if cp:
          pos += 1
          # if s1 is followed by *, then the monomial multiplying it follows
          if pos < ns and s[pos] == '*':
            pos += 1
            # pos1= position of next coefficient subexpression
            pos1 = s.find('+(',pos)
            # +(
            prev_pos = pos1+2
            if pos1 == -1:
              expv = self.read_monom(s[pos:])
            else:
              expv = self.read_monom(s[pos:pos1])
          else:
            expv = ''
            pos1 = -1
          p[expv] = cp
          if pos1 == -1 or pos >= ns:
            break
          else:
            pos = pos1 + 1
          if pos >= ns:
            break
        else:
          # the coefficient is zero, go to the next term
          if pos < ns:
            pos1 = s.find('+(',pos)
            if pos1 == -1:
              break
            prev_pos = pos1+2
            pos = prev_pos
      else:
        pos += 1
        if pos >= ns:
          break

    if nparens != 0:
      raise ValueError, 'nparens=%d' % nparens
    return p




  def __call__(self,s):
    """generate a free algebra element
    s can be a FreeAlgebraElement, belong to its coefficient ring,
    or be a string of the form
    c_0*m_0 + c_1*m_1 + ...
    where c_i are coefficients and m_i are monomials

    >>> from rmpoly import *
    >>> from free_algebra import FreeAlgebra
    >>> from fractions import Fraction as mpq
    >>> fa,A,B = fgens('A,B',mpq)
    >>> fa('1/2 + 2/3*A*B^2*A + B^2')
     +B^2 +2/3*A*B^2*A +1/2
    """
    if isinstance(s, FreeAlgebraElement):
      if s.fa == self:
        return s
    p = FreeAlgebraElement(self)
    if not isinstance(s, str):
      try:
        c = self.field(s)
        if c:
          p[''] = c
        return p
      except:
        raise ValueError, 's cannot be converted to self.field'
    if self.parens:
      return self.read(s)
    s = s.strip()
    a = []
    if (s[0] == '+' or s[0] == '-'):
      sgn = s[0]
    else:
      sgn = '+'
      s = '+' + s
    it = rpm.finditer(s)
    it.next()
    s1 = s[1:]
    if rpm.search(s1):
      start = 1
      for match in it:
        t=match.span()
        a.append((sgn,s[start:t[0]]))
        start=t[0]+1
        sgn = s[t[0]]
      a.append((sgn, s[t[1]:]))
    else:
      a.append((sgn,s1))
    p = FreeAlgebraElement(self)
    for sgn, s1 in a:
      m = self.mon_eval(s1)
      if sgn == '-':
        m = (m[0],-m[1])
      if m[0] in p:
        p[m[0]] += m[1]
      else:
        p[m[0]] = m[1]
    p.simplify()
    return p

  def add_slice_rules(self,*t):
    """
    replace monomials with expressions

    t = mon1,expr1, mon2,expr2, ...
    mon1,mon2,... are free algebra monomials
    expr1,expr2,... are free algebra expressions or 0.
    This rule searches for a monomial mon_i which is included
    in a monomial term in self and replaces it with expr_i

    >>> from rmpoly import *
    >>> from free_algebra import FreeAlgebra
    >>> from fractions import Fraction as mpq
    >>> fa,A,B = fgens('A,B',mpq)
    >>> fa.add_slice_rules(B*A,A*B+1)
    >>> fa.set_rules()
    >>> 7*A*B*A + 11*A*B*A**3
     +7*A^2*B +11*A^4*B +33*A^3 +7*A
    """
    if len(t)%2 != 0:
      raise ValueError, 't[2*i], t[2*i+1] define a rule'
    a = []
    for i in range(0,len(t),2):
      keys = t[i].keys()
      if len(keys) != 1:
        raise NotImplementedError, 'lhs must be a monomial in the free algebra generators'
      lhs = keys[0]
      rhs = t[i+1]
      if not isinstance(rhs, FreeAlgebraElement):
        #rhs = ((),self.field(rhs))
        rhs = self(rhs)
      else:
        rhs = rhs.items()
      a.append((lhs,rhs))
    self._pre_rules.append((slice_rule,tuple(a)))

  def set_rules(self):
    self.rules = tuple(self._pre_rules)

  def clear_rules(self):
    self.rules = ()

class FreeAlgebra(BaseFreeAlgebra):
  def __init__(self, pol_gens, field, **kwds):
    BaseFreeAlgebra.__init__(self, pol_gens, field, **kwds)
    self.__name__ = 'FreeAlgebra'

class MFreeAlgebra(BaseFreeAlgebra):
  """ class of free algebra on a ring with elements with more than one term
  """
  def __init__(self, pol_gens, field, **kwds):
    BaseFreeAlgebra.__init__(self, pol_gens, field, **kwds)
    self.parens = True
    self.__name__ = 'MFreeAlgebra'

def fgens(pol_gens, field, **kwds):
  """
  factory function to generate FreeAlgebra object and its generators
  """
  fa = FreeAlgebra(pol_gens, field, **kwds)
  if fa.parens:
    fa.__class__ == MFreeAlgebra
    fa.__name__ = 'MFreeAlgebra'
  else:
    fa.__class__ == FreeAlgebra
    fa.__name__ = 'FreeAlgebra'
  return [fa] + fa.gens()

def mfgens(pol_gens, field, **kwds):
  """
  factory function to generate MFreeAlgebra object and its generators
  """
  fa = MFreeAlgebra(pol_gens, field, **kwds)
  fa.__class__ == MFreeAlgebra
  fa.__name__ = 'MFreeAlgebra'
  return [fa] + fa.gens()

class FreeAlgebraElement(dict):
  """
  let A, B = fa.gens()
  a monomial in the fae, e.g. 3*A^4*B^3*A^2, is represented by the item
  (expv,v) with expv = (0,0,0,0,1,1,1,0,0) and v = 3

  """
  def __init__(self,fa,**kw):
    self.fa = fa
    dict.__init__(self,**kw)

  def copy(self):
    return copy(self)

  def __str__(self):
    #print 'DB30', dict(self)
    if not self:
      return '0'
    if self.fa.parens:
      return self.tostr()
    s = ''
    a = self.keys()
    a.sort(reverse=True)
    dpol_gens = self.fa.dpol_gens
    #print 'DB30a a=',a
    for expv in a:
      sa = []
      c = self[expv]
      # b list of [e,pw]
      b = []
      if expv:
        prev = expv[0]
        b.append([expv[0],1])
        for i in range(1,len(expv)):
          if expv[i] == prev:
            b[-1][1] += 1
          else:
            b.append([expv[i],1])
          prev = expv[i]
      #print 'DB31', expv, c, c > 0
      if c > 0:
        #print 'DB31+'
        s += ' +'
      else:
        s += ' -'
      if c < 0:
        c = -c
      if c != 1 or not expv:
        cnt1 = str(c)
      else:
        cnt1 = ''
      for i,x in enumerate(b):
        if x[1] == 1:
          sa.append(dpol_gens[x[0]])
        else:
          sa.append('%s^%d' % (dpol_gens[x[0]],x[1]))
      if cnt1:
        sa = [cnt1] + sa
      #print 'DB32 cnt1=%s, sa=%s, s=%s' %(cnt1,sa, s)
      s += '*'.join(sa)
      #print 'DB33', s
    return s
    #return str(dict(self))


  def tostr(self):
    fa = self.fa
    s = ''
    a = self.keys()
    a.sort(reverse=True)
    dpol_gens = self.fa.dpol_gens
    s = ''
    for expv in a:
      c = self[expv]
      s += ' +(%s)' % c
      b = []
      if expv:
        prev = expv[0]
        b.append([expv[0],1])
        for i in range(1,len(expv)):
          if expv[i] == prev:
            b[-1][1] += 1
          else:
            b.append([expv[i],1])
          prev = expv[i]

      if expv:
        s += '*'
        for i,x in enumerate(b):
          if x[1] == 1:
            s += dpol_gens[x[0]]
          else:
            s += '%s^%d' % (dpol_gens[x[0]],x[1])
    return s

  __repr__ = __str__

  def strip(p):
    """eliminate monomials with zero coefficient
    """
    z = p.fa.field(0)
    for k,v in p.items():
      if v == z:
        del p[k]

  def simplify(self):
    """apply rules if present, then eliminate monomials with zero coefficient
    
    >>> from rmpoly import *
    >>> from gmpy import mpq
    >>> from free_algebra import *
    >>> fa,A,Ad,Bra,Ket = fgens('A,Ad,Bra,Ket',mpq)
    >>> fa.add_slice_rules(A*Ket,0, Bra*Ad,0, Bra*Ket,fa(1), A*Ad,Ad*A+1)
    >>> fa.set_rules()
    >>> (A + Ad)**2
     +Ad^2 +2*Ad*A +A^2 +1
    >>> (A + Ad)**2*Ket
     +Ket +Ad^2*Ket
    >>> Bra*(A + Ad)**2*Ket
     +1
    >>> Bra*(A + Ad)**4*Ket
     +3
    """
    rules = self.fa.rules
    unchanged = {}
    if rules:
      while 1:
        changed=0
        keys = self.keys()
        for expv in keys:
          if not expv or expv in unchanged:
            continue
          for xx in rules:
            rule, d = xx
            ch = rule(self,expv,d)
            if not ch:
              unchanged[expv] = 1
            changed += ch
        if not changed:
          break
    self.strip()

  def __add__(p1,p2):
    if not p2:
      return p1.copy()
    if isinstance(p2,FreeAlgebraElement):
      if p1.fa == p2.fa:
        p = FreeAlgebraElement(p1.fa)
        for k,v in p1.iteritems():
          if k in p2:
            r = v + p2[k]
            if r:
              p[k] = r
          else:
            p[k] = v
        for k,v in p2.iteritems():
          if k not in p1:
            p[k] = v
      else:
        raise NotImplementedError
    elif isinstance(p2,Poly):
        return p2 + p1
    else:
      p = p1.copy()
      if '' in p1:
        p[()] = p[()] + p2
      else:
        p[''] = p2
    p.simplify()
    return p

  def __radd__(p1,n):
    p = p1.copy()
    if not n:
      return p
    if '' not in p1.keys():
      p[''] = p1.fa.field(n)
    else:
      if n == -p['']:
        del p['']
      else:
        p[''] += n
    p.simplify()
    return p

  def __neg__(self):
    return self*(-1)
  
  def __pos__(self):
    return self

  def __sub__(p1,p2):
    if not p2:
      return p1.copy()
    if isinstance(p2,FreeAlgebraElement):
      if p1.fa == p2.fa:
        p = FreeAlgebraElement(p1.fa)
        for k in p1:
          if k in p2:
            r = p1[k] - p2[k]
            if r:
              p[k] = r
          else:
            p[k] = p1[k]
        for k in p2:
          if k not in p1:
            p[k] = -p2[k]
    elif isinstance(p2,Poly):
      p = p2 - p1
      p = -p
    p.strip()
    return p

  def __mul__(p1,p2):
    fa = p1.fa
    p = FreeAlgebraElement(fa)
    if not p2:
      return p
    if isinstance(p2,FreeAlgebraElement):
      if fa == p2.fa:
        get = p.get
        p2it = p2.items()
        for exp1,v1 in p1.iteritems():
          for exp2,v2 in p2it:
            exp = exp1 + exp2
            v = v1*v2
            p[exp] = get(exp,0) + v1*v2
    elif isinstance(p2,Poly):
      return p2.__rmul__(p1)
    else:
      try:
        p2 = fa.field(p2)
        for exp1,v1 in p1.iteritems():
          p[exp1] = v1*p2
      except:
        raise NotImplementedError
    p.simplify()
    return p

  def __rmul__(p1,p2):
    p = FreeAlgebraElement(p1.fa)
    if not isinstance(p2,FreeAlgebraElement):
      if not p2:
        return p
      for exp1,v1 in p1.iteritems():
        p[exp1] = p2*v1
    p.simplify()
    return p

  def square(p1):
    # TODO optimize
    p = p1*p1
    # useless to call simpify
    return p

  def __pow__(self,n):
    fa = self.fa
    if not hasattr(n,'__hex__'):
      raise NotImplementedError
    n = int(n)
    if n < 0:
      raise ValueError, 'negative powers not allowed'
    if n == 0:
      if self:
        return fa(1)
      else:
        raise ValueError, 'Error 0^0'
    elif len(self) == 1:
      p = FreeAlgebraElement(fa)
      k, v = self.items()[0]
      p[k*n] = v**n
      p.simplify()
      return p
    elif n == 1:
      p = self.copy()
    elif n == 2:
      p = self.square()
    elif n == 3:
      p = self*self.square()
      return p
    p = fa(1)
    while 1:
      if n&1:
        p = p*self
        n -= 1
        if not n:
          break
      self = self.square()
      n = n // 2

    # useless to call simpify
    return p

  def __eq__(p1,p2):
    if not p2 or p2 == 0:
      return not p1
    if isinstance(p2,FreeAlgebraElement):
      if p1.fa == p2.fa:
        return dict.__eq__(p1,p2)
    else:
      if '' not in p1 or len(p1) > 1:
        return False
      return p1[''] == p1.fa.field(p2)

  def __div__(p1,n):
    p = FreeAlgebraElement(p1.fa)
    if not n:
      raise ZeroDivisionError
    for exp1,v in p1.iteritems():
      p[exp1] = v/n
    p.simplify()
    return p

  def coeff(self,p1):
    """monomial coefficient

    >>> from fractions import Fraction as mpq
    >>> fa = FreeAlgebra('A,B',mpq)
    >>> A,B = fa.gens()
    >>> c = (1 + A + B + A*B)**2
    >>> c.coeff(A*B) == 3
    True
    """
    if len(p1) != 1:
      raise ValueError, "p1 must be a monomial"
    expv = p1.keys()[0]
    if expv in self:
      return self[expv]
    else:
      return self.fa(0)

def slice_rule(self,expv,d):
  """
  expv is a key of self
  d dictionary with the rule
  if there is successful application of the rule, return True;
  else return False
  example:
  A gen. 0; B gen. 1
  rule B*A = A*B + 1
  The LHS has expv1 = (1,0); see in expv if there is a subtuple
  (slice) equal to expv1, to substitute
  For each slice rule there is an item in d
  d={(1,0):(((0, 1), mpq(1)), ((), mpq(1)))}
  """
  v = self[expv]
  for e,t in d:
    if e == expv:
      del self[expv]
      for e1,v1 in t:
        if e1 in self:
          self[e1] += v1*v
        else:
          self[e1] = v1*v
      return True
    else:
      n = len(e)
      m = len(expv)
      if m <= n:
        continue
      i = expv.find(e)
      if i == -1:
        continue
      del self[expv]
      for e1, v1 in t:
        if not v1:
          continue
        e1 = expv.replace(e,e1,1)
        if e1 in self:
          self[e1] += v1*v
        else:
          self[e1] = v1*v
      return True
  return False

if __name__ == "__main__":
    import doctest
    import sys
    if sys.version_info < (2, 6):
      print 'doctests require Fraction, available from Python2.6'
      sys.exit()
    doctest.testmod()

